{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "\n",
    "print(f\"Working directory: {os.getcwd()}\")\n",
    "# experiment_path = \"experiments\\\\2024-11-04_15-34-38\\\\9fba3d9_dr-sac_algorithm_environment_zero-init\\\\sac_l2f_false\"\n",
    "experiment_path = \"experiments\\\\2024-11-04_16-29-38\\\\9fba3d9_dr-sac_algorithm_environment_zero-init\\\\sac_l2f_false\"\n",
    "print(f\"Loading returns from {experiment_path}\")\n",
    "\n",
    "returns = []\n",
    "for seed in sorted(os.listdir(experiment_path)):\n",
    "    return_path_set = os.path.join(experiment_path, seed, \"return.json.set\")\n",
    "    return_path = os.path.join(experiment_path, seed, \"return.json\")\n",
    "    if os.path.exists(return_path_set):\n",
    "        with open(return_path, \"r\") as f:\n",
    "            returns.append(json.load(f))\n",
    "print(f\"Loaded {len(returns)} returns\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist([r[\"parameters\"][\"dynamics\"][\"mass\"] for r in returns])\n",
    "plt.ylabel(\"Count\")\n",
    "plt.xlabel(\"Mass [kg]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "thrusts = np.array([np.dot(r[\"parameters\"][\"dynamics\"][\"rotors\"][0][\"thrust_curve\"], [1, 1, 1]) for r in returns])\n",
    "weights = np.array([r[\"parameters\"][\"dynamics\"][\"mass\"] * 9.81 for r in returns])\n",
    "thrust_to_weight = thrusts * 4 / weights\n",
    "\n",
    "plt.hist(thrust_to_weight)\n",
    "plt.ylabel(\"Count\")\n",
    "plt.xlabel(\"Thrust to Weight Ratio\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inertia_xx = np.array([r[\"parameters\"][\"dynamics\"][\"J\"][0][0] for r in returns])\n",
    "plt.hist(inertia_xx)\n",
    "plt.ylabel(\"Count\")\n",
    "plt.xlabel(\"Inertia (xx) [kg*m^2]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rotor_distances = np.array([np.linalg.norm(r[\"parameters\"][\"dynamics\"][\"rotors\"][0][\"pose\"][\"position\"]) for r in returns])\n",
    "\n",
    "plt.hist(rotor_distances * 2)\n",
    "plt.ylabel(\"Count\")\n",
    "plt.xlabel(\"Diameter [m]\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torques = rotor_distances * np.sqrt(2) * thrusts * 9.81\n",
    "torque_to_inertia = torques / inertia_xx\n",
    "plt.hist(torque_to_inertia)\n",
    "plt.ylabel(\"Count\")\n",
    "plt.xlabel(\"Torque to Inertia Ratio\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(thrust_to_weight / torque_to_inertia)\n",
    "plt.ylabel(\"Count\")\n",
    "plt.xlabel(\"Thrust to Weight to Torque to Inertia Ratio\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_episode_lengths = [max([max(step[\"episode_length\"]) for step in r[\"evaluation\"]]) for r in returns]\n",
    "\n",
    "plt.hist(max_episode_lengths)\n",
    "plt.ylabel(\"Count\")\n",
    "plt.xlabel(\"Max Episode Lengths [steps]\")\n",
    "plt.show()\n",
    "\n",
    "mean_episode_lengths = [np.mean([np.mean(step[\"episode_length\"]) for step in r[\"evaluation\"]]) for r in returns]\n",
    "\n",
    "plt.hist(mean_episode_lengths)\n",
    "plt.ylabel(\"Count\")\n",
    "plt.xlabel(\"Mean Episode Lengths [steps]\")\n",
    "plt.show()\n",
    "\n",
    "final_mean_episode_length = [np.mean(r[\"evaluation\"][-1][\"episode_length\"]) for r in returns]\n",
    "plt.hist(final_mean_episode_length)\n",
    "plt.ylabel(\"Count\")\n",
    "plt.xlabel(\"Final Mean Lengths [steps]\")\n",
    "plt.show()\n",
    "\n",
    "steps = [step[\"step\"] for step in returns[0][\"evaluation\"]]\n",
    "\n",
    "learning_curves = [[np.mean(step[\"episode_length\"]) for step in r[\"evaluation\"]] for r in returns]\n",
    "\n",
    "for el in learning_curves:\n",
    "    plt.plot(steps, el)\n",
    "plt.ylabel(\"Count\")\n",
    "plt.xlabel(\"Max Episode Lengths [steps]\")\n",
    "plt.show()\n",
    "\n",
    "learning_curves_aggregate = np.stack(learning_curves, axis=0)\n",
    "plt.fill_between(steps, np.mean(learning_curves_aggregate, axis=0) - np.std(learning_curves_aggregate, axis=0), np.mean(learning_curves_aggregate, axis=0) + np.std(learning_curves_aggregate, axis=0), alpha=0.5)\n",
    "plt.plot(steps, np.mean(learning_curves_aggregate, axis=0))\n",
    "plt.ylabel(\"Episode Length\")\n",
    "plt.xlabel(\"Step\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "data = {\n",
    "    'mass': [r[\"parameters\"][\"dynamics\"][\"mass\"] for r in returns],\n",
    "    'rotor_distance': rotor_distances,\n",
    "    'inertia_xx': inertia_xx,\n",
    "    'thrust_to_weight': thrust_to_weight,\n",
    "    'torque_to_inertia': torque_to_inertia,\n",
    "    'thrust_to_weight_to_torque_to_inertia': thrust_to_weight / torque_to_inertia,\n",
    "    'max_episode_length': max_episode_lengths,\n",
    "    'mean_episode_length': mean_episode_lengths,\n",
    "    'final_mean_episode_length': final_mean_episode_length\n",
    "}\n",
    "df = pd.DataFrame(data)\n",
    "\n",
    "corr_matrix = df.corr()\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 6))\n",
    "cax = ax.matshow(corr_matrix, cmap='coolwarm')\n",
    "\n",
    "fig.colorbar(cax)\n",
    "\n",
    "ax.set_xticks(np.arange(len(corr_matrix.columns)))\n",
    "ax.set_yticks(np.arange(len(corr_matrix.columns)))\n",
    "ax.set_xticklabels(corr_matrix.columns, rotation=45, ha='left')\n",
    "ax.set_yticklabels(corr_matrix.columns)\n",
    "\n",
    "for i in range(len(corr_matrix.columns)):\n",
    "    for j in range(len(corr_matrix.columns)):\n",
    "        text = ax.text(j, i, f'{corr_matrix.iloc[i, j]:.2f}', ha='center', va='center', color='black')\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
